#include "pin.H"

#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <iostream>
#include <fstream>
#include <string>
#include <string.h>
#include <time.h>
#include <assert.h>
#include <vector>
#ifndef _WIN32
#include <sys/time.h>
#endif

#include "LoadPlugin.h"
#include "HSGlobal.h"
#include "PowerModels/PowerModel.h"
#include "PowerModels/ModelFactory.h"
#include "Devices/DeviceFactory.h"
#include "Devices/Device.h"
#include "Scheduler/Scheduler.h"
#include "Checkpointer/PinCheckpointer.h"
#include "Runtime/HarvSimRuntime_Pin.h"

std::vector<Scheduler *> Schedulers;

KNOB<string> KnobPowerModel(KNOB_MODE_WRITEONCE, "pintool",
			   "models", "Plugins/UniformPowerModel.so", "Power Model Modules To Simulate");

KNOB<string> KnobDevices(KNOB_MODE_WRITEONCE, "pintool",
			   "devices", "Devices/RandomSensor.so", "External Devices To Simulate");

KNOB<string> KnobRecoveryRoutines(KNOB_MODE_WRITEONCE, "pintool",
			   "recovery", "", "Paths to recovery routine plugins - The order should correspond with the recovery block's numeric name.");

FILE *outfile;
KNOB<string> KnobOutFile(KNOB_MODE_WRITEONCE, "pintool",
			   "outfile", "harvsim.out", "Output file to report to. Specifying \"-\" will report stats to stdout");

#define MAX_NTHREADS 64
unsigned long instrumentationStatus[MAX_NTHREADS];

enum MemOpType { MemRead = 0, MemWrite = 1 };

INT32 usage()
{
    cerr << "-=HarvSim=- An interruption-prone device simulator";
    cerr << KNOB_BASE::StringKnobSummary();
    cerr << endl;
    return -1;
}


VOID TurnInstrumentationOn(ADDRINT tid){
  instrumentationStatus[PIN_ThreadId()] = true; 
}

VOID TurnInstrumentationOff(ADDRINT tid){
  instrumentationStatus[PIN_ThreadId()] = false; 
}

VOID instrumentRoutine(RTN rtn, VOID *v){
 
    
  if(strstr(RTN_Name(rtn).c_str(),"INSTRUMENT_OFF")){
    RTN_Open(rtn);
    RTN_InsertCall(rtn, 
                   IPOINT_BEFORE, 
                   (AFUNPTR)TurnInstrumentationOff, 
                   IARG_THREAD_ID,
                   IARG_END);
    RTN_Close(rtn);
  }
   
  
  if(strstr(RTN_Name(rtn).c_str(),"INSTRUMENT_ON")){
    RTN_Open(rtn);
    RTN_InsertCall(rtn, 
                   IPOINT_BEFORE, 
                   (AFUNPTR)TurnInstrumentationOn, 
                   IARG_THREAD_ID,
                   IARG_END);
    RTN_Close(rtn);
  }

}

VOID instrumentImage(IMG img, VOID *v)
{

}

void Read(ADDRINT inst, ADDRINT opea){


  std::vector<Scheduler *>::iterator i,e;
  for(i = Schedulers.begin(), e = Schedulers.end(); i != e; i++){

    (*i)->scheduleInstruction( inst, (*i)->readCost(inst,opea) );

  }

}

void Write(ADDRINT inst, ADDRINT opea){

  std::vector<Scheduler *>::iterator i,e;
  for(i = Schedulers.begin(), e = Schedulers.end(); i != e; i++){

    (*i)->scheduleInstruction( inst, (*i)->writeCost(inst,opea) );

  }

}

void Inst(ADDRINT inst, INT32 category){

  std::vector<Scheduler *>::iterator i,e;
  for(i = Schedulers.begin(), e = Schedulers.end(); i != e; i++){

    (*i)->scheduleInstruction( inst, (*i)->insCost(category) );

  }

}

VOID instrumentTrace(TRACE trace, VOID *v)
{

  for (BBL bbl = TRACE_BblHead(trace); BBL_Valid(bbl); bbl = BBL_Next(bbl)) {

    for (INS ins = BBL_InsHead(bbl); INS_Valid(ins); ins = INS_Next(ins)) {  

      if(INS_IsMemoryRead(ins)) {

	  INS_InsertCall(ins, 
			 IPOINT_BEFORE, 
			 (AFUNPTR)Read, 
			 IARG_INST_PTR,
			 IARG_MEMORYREAD_EA,
			 IARG_END);

      } else if(INS_IsMemoryWrite(ins)) {

	  INS_InsertCall(ins, 
			 IPOINT_BEFORE, 
			 (AFUNPTR)Write, 
			 IARG_INST_PTR,//instruction address of write
			 IARG_MEMORYWRITE_EA,//address being accessed
			 IARG_END);

      }else{

	  INS_InsertCall(ins, 
			 IPOINT_BEFORE, 
			 (AFUNPTR)Inst, 
			 IARG_INST_PTR,
                         IARG_UINT32, INS_Category( ins ),
			 IARG_END);

      }

    }

  }

}


VOID threadBegin(THREADID threadid, CONTEXT *sp, INT32 flags, VOID *v)
{
  
}
    
VOID threadEnd(THREADID threadid, const CONTEXT *sp, INT32 flags, VOID *v)
{

}

VOID dumpInfo(){
}


VOID Fini(INT32 code, VOID *v)
{

#ifndef _DINO_WINDOWS_ 
  struct timeval t;
  gettimeofday(&t,NULL);
  char *datestring = ctime(&(t.tv_sec));
  HSLOG(outfile,"End of Simulation %s",datestring);
#else
  HSLOG(outfile,"End of Simulation\n");
#endif
  HSLOG(outfile,"Statistics:\n");
  HSLOG(outfile,"-----------------------\n");
  std::vector<Scheduler *>::iterator i,e;
  for(i = Schedulers.begin(), e = Schedulers.end(); i != e; i++){

    (*i)->dumpStats(outfile);

  }
  
}

BOOL segvHandler(THREADID threadid,INT32 sig,CONTEXT *ctx,BOOL hasHndlr,const EXCEPTION_INFO *pExceptInfo, VOID*v){
  return TRUE;//let the program's handler run too
}

BOOL termHandler(THREADID threadid,INT32 sig,CONTEXT *ctx,BOOL hasHndlr,const EXCEPTION_INFO *pExceptInfo, VOID*v){
  return TRUE;//let the program's handler run too
}


int main(int argc, char *argv[])
{

  PIN_InitSymbols();
  if( PIN_Init(argc,argv) ) {
    return usage();
  }

  
  const char *outfilestr = KnobOutFile.Value().c_str();
  if(!strncmp(outfilestr,"-",1)){

    outfile = stdout;

  }else{

    outfile = fopen(outfilestr,"w");

  }
  
#ifndef _DINO_WINDOWS_ 
  struct timeval t;
  gettimeofday(&t,NULL);
  char *datestring = ctime(&(t.tv_sec));
  HSLOG(outfile,"A Simple Energy Harvesting Model Simulator %s",datestring);
#else
  HSLOG(outfile,"A Simple Energy Harvesting Model Simulator\n");
#endif

  const char *pstr = KnobPowerModel.Value().c_str();
  char *ct = strtok((char *)pstr,",");
  while(ct != NULL){

    HSLOG(outfile,"Opening power model plugin: \"%s\"\n",ct);
    void *chand = OpenLibrary( ct );
    if( chand == NULL ){
      HSLOG(outfile,"Couldn't Load %s\n", ct);
      //HSLOG(outfile,"dlerror: %s\n", dlerror());
      exit(1);
    }

    //dlerror(); 
    ModelFactory cfac = (ModelFactory)GetLibraryFuncByName(chand, (char*)"Create");

    if( cfac == NULL ){

      HSLOG(outfile,"Couldn't get the Create function\n");
      //HSLOG(outfile,"dlerror: %s\n", dlerror());
      exit(1);

    }

    PowerModel *m = cfac();
    Stash *s = new Stash(16384);
    
    Schedulers.push_back( new Scheduler(m,s) );
    
    HSLOG(outfile,"Created a new scheduler with the \"%s\"\n",m->Identify());

    ct = strtok(NULL,","); 

  }

  const char *dstr = KnobDevices.Value().c_str();
  char *dt = strtok((char *)dstr,",");
  while(dt != NULL){

    HSLOG(outfile,"Opening device plugin: \"%s\"\n",dt);
    void *chand = OpenLibrary( dt );
    if( chand == NULL ){
      HSLOG(outfile,"Couldn't Load %s\n", dt);
      //HSLOG(outfile,"dlerror: %s\n", dlerror());
      exit(1);
    }

    //dlerror(); 
    DeviceFactory cfac = (DeviceFactory)GetLibraryFuncByName(chand, (char*)"Create");

    if( cfac == NULL ){

      HSLOG(outfile,"Couldn't get the Create function\n");
      //HSLOG(outfile,"dlerror: %s\n", dlerror());
      exit(1);

    }

    Device *d = cfac();
    HSLOG(outfile,"Created a new device: \"%s\"\n",d->Identify());

    dt = strtok(NULL,","); 

  }
 
  vector< void(*)(Recovery *) > recos; 
  recos.clear();
  const char *rstr = KnobRecoveryRoutines.Value().c_str();
  char *rt = strtok((char *)rstr,",");
  while(rt != NULL){

    HSLOG(outfile,"Opening recovery routine plugin: \"%s\"\n",rt);
    void *chand = OpenLibrary( rt );
    if( chand == NULL ){
      HSLOG(outfile,"Couldn't Load %s\n", rt);
      //HSLOG(outfile,"dlerror: %s\n", dlerror());
      exit(1);
    }

    //dlerror(); 
    void (*r)(Recovery *) = (void (*)(Recovery *))GetLibraryFuncByName(chand, (char*)"Recover");

    if( r == NULL ){

      HSLOG(outfile,"Couldn't get the Create function\n");
      //HSLOG(outfile,"dlerror: %s\n", dlerror());
      exit(1);

    }
    HSLOG(outfile,"Opened recovery routine plugin: \"%s\"\n",rt);
    recos.push_back( r );

    rt = strtok(NULL,","); 

  }



  RTN_AddInstrumentFunction(instrumentRoutine,0);
  IMG_AddInstrumentFunction(instrumentImage, 0);
  TRACE_AddInstrumentFunction(instrumentTrace, 0);

  initRuntime();
  initCheckpointer(&recos);

  PIN_InterceptSignal(SIGTERM,termHandler,0);
  PIN_InterceptSignal(SIGSEGV,segvHandler,0);

  PIN_AddThreadStartFunction(threadBegin, 0);
  PIN_AddThreadFiniFunction(threadEnd, 0);
  PIN_AddFiniFunction(Fini, 0);
    
 
  PIN_StartProgram();
  
  return 0;
}
